from langchain.prompts import PromptTemplate
from pydantic import BaseModel, Field
from datetime import datetime
import json
from LLM_Prompt.siliconflow import SiliconFlow

# 使用 SiliconFlow 或 OpenAI 接入大模型
llm = SiliconFlow()  # 替换成你实际的 SiliconFlow 模型接口

classification_prompt = PromptTemplate(
    input_variables=["question", "current_date"],
    template="""  
# 使用大模型进行分类与参数提取  

## 可用 API  
1. `vessel_trajectory`:  
   - **用途**：获取船只的航线轨迹数据。  
   - **必需参数**：  
     - `boatName`（船只名称）  
     - `startTime`（查询起始时间，ISO8601 格式）  
     - `endTime`（查询结束时间，ISO8601 格式）  

2. `checkpoint_passage`:  
   - **用途**：监控船只经过检查点的情况。  
   - **必需参数**：  
     - `boatName`（船只名称）  
     - `startTime`（查询起始时间，ISO8601 格式）  
     - `endTime`（查询结束时间，ISO8601 格式）  

3. `violation_detection`:  
   - **用途**：查询船只的违规行为。  
   - **必需参数**：  
     - `boatName`（船只名称）  
     - `startTime`（查询起始时间，ISO8601 格式）  
     - `endTime`（查询结束时间，ISO8601 格式）  

4. `checkpoint_traffic`:  
   - **用途**：查询特定卡口通过的船舶信息。  
   - **必需参数**：  
     - `checkpointName`（卡口名称）  
     - `startTime`（查询起始时间，ISO8601 格式）  
     - `endTime`（查询结束时间，ISO8601 格式）  

5. `current_waterway_status`:  
   - **用途**：查询特定航道或港区当前的船舶情况。  
   - **必需参数**：  
     - `waterwayName`（航道名称）  

6. `not matched`:  
   - **用途**：以上API均未匹配上则进入此API。  
   - **必需参数**：  

## 模型任务  
1. 根据输入问题分类，从上述 API 中选择一个最相关的 API.
2. 今天日期是{current_date}提取相关的参数，包括：  
   - `boatName`：船只名称（如适用）。  
   - `startTime`：查询起始时间（如适用，ISO8601 格式，如 `2025-04-09 00:00:00`）。  
   - `endTime`：查询结束时间（如适用，ISO8601 格式，如 `2025-04-09 23:59:59`）。  
   - `checkpointName`：卡口名称（如适用）。  
   - `waterwayName`：航道名称（如适用）。  
3. 未提及时间则返回 \"未知\",相对时间（如“今天”、“昨天”或“前天”）需根据当前日期 {current_date} 计算
4. 卡口名称从以下内容全匹配，或者模糊匹配：杭申线卡口，长湖申线卡口，金汇港(北闸)，油墩港(油墩港船闸)，苏申内港线(上港宜东)，吴淞江(泗江口站)，黄浦江(松浦大桥南岸)，龙泉港(叶榭塘船闸)，
油墩港(东大盈船闸)，大治河船闸，苏申外港线省际检查站，杨思船闸口，东沟船闸口，苏申内港线省际检查站，平申线省际检查站

## 输出格式  
根据分类与参数提取结果，返回以下 JSON 格式：  
```json  
{{  
  "api": "API 名称",  
  "parameters": {{  
    "boatName": "示例船只名称或 null",  
    "startTime": "2025-04-08 00:00:00 或 null",  
    "endTime": "2025-04-08 23:59:59 或 null",  
    "checkpointName": "卡口名称或 null",  
    "waterwayName": "航道名称或 null"  
  }}  
}}  
示例
输入
问题：{question}

输出
返回严格按照上述 JSON 格式的分类与参数提取结果，如果有多个api符合只返回一个你觉得最符合的结果，不需要输出其他内容。
"""
)

# Step 4: Pydantic 数据验证器
class ExtractedData(BaseModel):
    boat_name: str = Field(..., description="船的名字")
    date: str = Field(..., description="日期，格式为 YYYY-MM-DD")

    @staticmethod
    def validate_output(response: str):
        try:
            data = ExtractedData.model_validate_json(response)  # 替换 parse_raw
            return data.model_dump()  # 替换 dict
        except Exception as e:
            raise ValueError(f"数据解析错误: {e}")


# 清理模型输出工具
def clean_model_output(raw_output: str) -> str:
    """
    清理大模型的输出，去除额外的 Markdown 等格式符号，确保输出是合法 JSON。
    """
    # 去掉 Markdown 格式的代码块标记（如 ```json ... ```)
    if raw_output.startswith("```json"):
        raw_output = raw_output.strip("```json").strip("`\n")
    elif raw_output.startswith("```"):
        raw_output = raw_output.strip("```").strip("`\n")

    try:
        # 确保最终输出为合法的 JSON 格式（如有多余的空格，再清理一下）
        json.loads(raw_output)  # 测试是否是有效 JSON
    except json.JSONDecodeError as e:
        raise ValueError(f"清理后的输出仍不是有效 JSON: {e}")

    return raw_output



# Step 5: 链式操作
classification_chain = classification_prompt | llm


def classify_and_extract(question):
    # 动态获取当前日期
    current_date = datetime.now().strftime("%Y-%m-%d")

    # 1. 问题分类 提取关键信息
    classification_extract_json = classification_chain.invoke({"question": question, "current_date":current_date})

    # 清理模型输出
    cleaned_result = clean_model_output(classification_extract_json)

    return cleaned_result


# 测试代码
test_queries = [
        "莲花号前天的航道轨迹",
        "莲花号2025年1月20号的航道轨迹",
        "船舶 ABC123 2025年2月20日的行驶轨迹？",
        "船舶 XYZ789 昨天经过那些卡口？",
        "船舶 DEF456 本周是否有违规行为（预警）？",
        "船舶 GHI789 上周去过哪些码头？去过哪些区域？途径那些船闸？途径那些航道？",
        "卡口 南京长江大桥 2025年2月20日到22日有哪些船舶经过？有哪些违规船舶？有哪些集装箱船舶经过等",
        "航道 长江主航道 当前有哪些船舶？有哪些违规船舶？有哪些类别船舶？",
        "2025年2月20日至22日集装箱船的行驶轨迹？",
        "昨天有哪些集装箱船经过南京长江大桥卡口？",
        "本周有没有集装箱船存在违规行为？",
        "长江主航道当前有哪些集装箱船？",
        "上海2025年1月10日天气怎么样",
        "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。"]

# for example_question in test_queries:
#   print('example_question:', example_question)
#   result = classify_and_extract(example_question)
#   print(result)

# question = "分析2025年1月油墩港(东大盈船闸)每天的流量组成"
# result = classify_and_extract(question)
# print(result)